###########################################################
#### figure_s4.R                                       ####
#### Generates Figure s4 and accompanying descriptives ####
###########################################################

tab_in = import(here("Data","var_data.csv"))
vars_fact = tab_in |> 
  filter(is_fact == 1) |> 
  distinct(var2) |> 
  pull(var2)

data_recoded = data |> 
  mutate(pid = case_when(
    pid == "Democrat" ~ 1,
    pid == "Republican" ~ 2,
    pid == "Independent" ~ 3
  )) |> 
  mutate(across(all_of(vars_fact), as.factor))

pred_probs_ri = map_dfr(unique(tab_in$var), \(x){
  var_in = filter(tab_in, var == x)
  # Make percentiles
  if(!is.na(unique(var_in$ptile))){
    pred_vals = data.frame(
      val_label = c("Low (25th Perc.)", "Median", "High (75th Perc.)"),
      val = quantile(data_recoded[[x]], c(.25,.5,.75), na.rm = T)
    )
  } else {
    pred_vals = distinct(var_in, val_label, val)
  }
  if(unique(var_in$is_fact == 1)){
    pred_vals = mutate(pred_vals, val = as.factor(val))
  }
  names(pred_vals)[2] = x
  # Model
  f = as.formula(paste0("affpol ~ ", x, " + (1|state)"))
  mod = lmer(f, weights = wt_adj, data = data_recoded)
  # Predict
  predict(mod, newdata = pred_vals, re.form = NA, type = 'response', se.fit = T) |> 
    as_tibble() |> 
    bind_cols(pred_vals) |> 
    select(fit:val_label) |> 
    left_join(var_in)
})

fig_s4 = pred_probs_ri |> 
  ggplot(aes(x = fit, y = str_wrap(reorder(var_label, fit), 15), group = val_label)) +
  facet_wrap(~ class, ncol = 4, scales = "free_y") +
  geom_vline(xintercept = 50, alpha = .25) +
  geom_pointrange(aes(xmin = fit - 1.96*se.fit,
                      xmax = fit + 1.96*se.fit),
                  position = position_dodge(0.65)) +
  geom_text(aes(label = val_label, x = fit + 1.96*se.fit + .25), hjust = 0, size = 3,
            position = position_dodge(0.65)) +
  labs(y = NULL, x = 'Predicted Affective Polarization') +
  theme_prl() +
  theme(axis.text.y = element_text(lineheight = 0.7))

print(fig_s4)

ggsave(here("Plots","figure_s4.pdf"), fig_s4,
       units = "in", height = 4, width = 14, dpi = 600)
